library(estimatr)
library(lmtest)
library(dplyr)
library(kableExtra)

path_to_data <- "data/first_survey_formatted.rds"
survey_data <- readRDS(path_to_data)

survey_data$ResponseId <- as.factor(survey_data$ResponseId)
survey_data$task_number <- as.factor(survey_data$task_number)


task_number_uni_model <- lm_robust(vote_choice ~ University*task_number, clusters = ResponseId, data = survey_data) 
uni_model <- lm_robust(vote_choice ~ University, clusters = ResponseId, data = survey_data) 

task_number_ocu_model <- lm_robust(vote_choice ~ Occupation*task_number, clusters = ResponseId, data = survey_data) 
ocu_model <- lm_robust(vote_choice ~ `Occupation`, clusters = ResponseId, data = survey_data) 

task_number_hs_model <- lm_robust(vote_choice ~ `High School`*task_number, clusters = ResponseId, data = survey_data) 
hs_model <- lm_robust(vote_choice ~ `High School`, clusters = ResponseId, data = survey_data) 

task_number_fj_model <- lm_robust(vote_choice ~ `Father's Job`*task_number, data = survey_data) 
fj_model <- lm_robust(vote_choice ~ `Father's Job`, clusters = ResponseId, data = survey_data) 

task_number_uni_model_region <- lm_robust(good_region ~ University*task_number, clusters = ResponseId, data = survey_data) 
uni_model_region <- lm_robust(good_region ~ University + task_number, clusters = ResponseId, data = survey_data) 

task_number_ocu_model_region <- lm_robust(good_region ~ Occupation*task_number, clusters = ResponseId, data = survey_data) 
ocu_model_region <- lm_robust(good_region ~ `Occupation` + task_number, clusters = ResponseId, data = survey_data) 

task_number_hs_model_region <- lm_robust(good_region ~ `High School`*task_number, clusters = ResponseId, data = survey_data) 
hs_model_region <- lm_robust(good_region ~ `High School` + task_number, clusters = ResponseId, data = survey_data) 

task_number_fj_model_region <- lm_robust(good_region ~ `Father's Job`*task_number, clusters = ResponseId, data = survey_data) 
fj_model_region <- lm_robust(good_region ~ `Father's Job` + task_number, clusters = ResponseId, data = survey_data) 


task_number_uni_model_country <- lm_robust(good_country ~ University*task_number, clusters = ResponseId, data = survey_data) 
uni_model_country <- lm_robust(good_country ~ `University` + task_number, clusters = ResponseId, data = survey_data) 

task_number_ocu_model_country <- lm_robust(good_country ~ Occupation*task_number, clusters = ResponseId, data = survey_data) 
ocu_model_country <- lm_robust(good_country ~ `Occupation` + task_number, clusters = ResponseId, data = survey_data) 

task_number_hs_model_country <- lm_robust(good_country ~ `High School`*task_number, clusters = ResponseId, data = survey_data) 
hs_model_country <- lm_robust(good_country ~ `High School` + task_number, clusters = ResponseId, data = survey_data) 

task_number_fj_model_country <- lm_robust(good_country ~ `Father's Job`*task_number, clusters = ResponseId, data = survey_data) 
fj_model_country <- lm_robust(good_country ~ `Father's Job` + task_number, clusters = ResponseId, data = survey_data) 


ftest_vc_uni <- waldtest(task_number_uni_model, uni_model, test = "F")
ftest_vc_ocu <- waldtest(task_number_ocu_model, ocu_model, test = "F")
ftest_vc_hs <- waldtest(task_number_hs_model, hs_model, test = "F")
ftest_vc_fj <- waldtest(task_number_fj_model, fj_model, test = "F")

ftest_gr_uni <- waldtest(task_number_uni_model_region, uni_model_region, test = "F")
ftest_gr_ocu <- waldtest(task_number_ocu_model_region, ocu_model_region, test = "F")
ftest_gr_hs <- waldtest(task_number_hs_model_region, hs_model_region, test = "F")
ftest_gr_fj <- waldtest(task_number_fj_model_region, fj_model_region, test = "F")

ftest_gc_uni <- waldtest(task_number_uni_model_country, uni_model_country, test = "F")
ftest_gc_ocu <- waldtest(task_number_ocu_model_country, ocu_model_country, test = "F")
ftest_gc_hs <- waldtest(task_number_hs_model_country, hs_model_country, test = "F")
ftest_gc_fj <- waldtest(task_number_fj_model_country, fj_model_country, test = "F")

ftests <- list(
  ftest_vc_uni,
  ftest_vc_ocu,
  ftest_vc_hs,
  ftest_vc_fj,
  ftest_gr_uni,
  ftest_gr_ocu,
  ftest_gr_hs,
  ftest_gr_fj,
  ftest_gc_uni,
  ftest_gc_ocu,
  ftest_gc_hs,
  ftest_gc_fj
  
)

fstats <- sapply(ftests, function(x) x$F[2])
pvals <- sapply(ftests, function(x) x$`Pr(>F)`[2])

fstats <- round(fstats, digits = 2)
pvals <- round(pvals, digits = 2)

fstats_table_content <- paste0(fstats, " (p=" , pvals, ")")

outcome <- c("Vote choice", "Regional Welfare", "National Welfare")
uni <- fstats_table_content[c(1, 5, 9)]
ocu <- fstats_table_content[c(2, 6, 10)]
hs  <- fstats_table_content[c(3, 7, 11)]
fj <- fstats_table_content[c(4, 8, 12)]

fstat_table <- as.data.frame(cbind(outcome, uni, ocu, hs, fj))
colnames(fstat_table) <- c("Outcome Variable", "University F Stat", "Occupation F Stat", "High School F Stat", "Father's Job F Stat")

print(kableExtra::kable(fstat_table, format = "latex"))

